# This Python file uses the following encoding: utf-8
############################################################################
# Copyright (c) 2025 Li Auto Inc. and its affiliates
# Licensed under the Apache License, Version 2.0(the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
############################################################################
from typing import Dict, Union
from enum import Enum, unique


@unique
class VersionStr(Enum):
    VCOS_100 = "vcos 1.0"
    VCOS_200 = "2.0"


@unique
class VersionNumber(Enum):
    """版本编号类型"""

    VCOS_100 = 1
    VCOS_200 = 2


VERSION_MAP = {
    VersionStr.VCOS_100.value: VersionNumber.VCOS_100.value,
    VersionStr.VCOS_200.value: VersionNumber.VCOS_200.value,
}


class TraceVCOS100Parser:
    VERSION_STR = VersionStr.VCOS_100.value
    VERSION_NO = VersionNumber.VCOS_100.value
    # 报文长度
    TRACE_IRQ_ENTRY_LEN = 24
    TRACE_IRQ_EXIT_LEN = 24
    TRACE_SCHED_WAKEUP_LEN = 20
    TRACE_SCHED_SWITCH_LEN = 24
    TRACE_REACTIVATE_LEN = 16
    TRACE_EVENT_START_LEN = 24
    TRACE_RUNNABLE_S_LEN = 20
    TRACE_FUNCTION_S_LEN = 40
    TRACE_SPINLOCK_LEN = 20
    TRACE_RESOURCE_LEN = 20
    TRACE_INTERRUPT_S_LEN = 16
    # 报文类型字段
    FTRACE_SCHED_SWITCH = 1
    FTRACE_WAKE_UP = 2
    FTRACE_IRQ_ENTRY = 3
    FTRACE_IRQ_EXIT = 4
    FTRACE_EVENT_START = 5
    FTRACE_RUNNABLE_START = 6
    FTRACE_RUNNABLE_END = 7
    FTRACE_FUNCTION_START = 8
    FTRACE_FUNCTION_END = 9
    FTRACE_GET_SPINLOCK = 10
    FTRACE_RELEASE_SPINLOCK = 11
    FTRACE_GET_RESOURCE = 12
    FTRACE_RELEASE_RESOURCE = 13
    FTRACE_DISABLE_ALL_INT = 14
    FTRACE_ENABLE_ALL_INT = 15
    FTRACE_SUSPEND_ALL_INT = 16
    FTRACE_RESUME_ALL_INT = 17
    FTRACE_SUSPEND_OS_INT = 18
    FTRACE_RESUME_OS_INT = 19
    FTRACE_SUSPEND_KERNEL_INT = 20
    FTRACE_RESUME_KERNEL_INT = 21
    FTRACE_DISABLE_GLOBAL = 22
    FTRACE_ENABLE_GLOBAL = 23
    FTRACE_TASK_REACTIVATE = 24
    # 状态值
    TASK_STATE_RUNNING = 0
    TASK_STATE_READY = 1
    TASK_STATE_BLOCKED = 2
    TASK_STATE_SUSPEND = 3
    TASK_STATE_PENDING = 4
    TASK_STATE_DELAYED = 5
    TASK_STATE_DELETED = 6

    def __init__(self):
        self.tick_to_ns: float = 0.0  # 1tick等于多少ns
        self.time_offset: int = 0  # 时间戳偏移(ns)

    def __getattr__(self, name):
        # 当访问不存在的属性时，返回 -1
        return -1

    def get_mapped_id(self, ori_id, core_id):
        """针对不同core的对象id进行映射"""
        return ori_id

    def get_trace_data_list(self, data_list: bytearray, added_data: bytearray = None):
        """将二进制数据分成trace数据帧"""
        trace_pack_start = bytearray([0xA5, 0xA5])
        trace_pack_end = bytearray([0x5A, 0x5A])
        trace_data_list: list = []
        total_len: int = len(data_list)
        cur_idx: int = 0
        message_cnt = 0
        while cur_idx < total_len:
            # 每个报文包至少长6字节
            if (cur_idx + 6) > total_len:
                return trace_data_list
            # 包的开始字符
            if data_list[cur_idx : cur_idx + 2] != trace_pack_start:
                cur_idx += 1
                continue
            # 包的长度
            pack_len: int = (data_list[cur_idx + 3] << 8) | data_list[cur_idx + 2]
            pack_end_idx: int = cur_idx + pack_len
            # 当前包无效，找下一个包
            if (
                (pack_end_idx > total_len)
                or (pack_len < 6)
                or (data_list[pack_end_idx - 2 : pack_end_idx] != trace_pack_end)
            ):
                cur_idx += 1
                continue
            # 当前包有效，获取每个数据帧
            message_idx: int = cur_idx + 4
            message_end_idx: int = pack_end_idx - 2
            while message_idx < message_end_idx:
                trace_data_length = data_list[message_idx]
                if (message_idx + 2) > message_end_idx:
                    print(f"[WARNING] parse trace data message_idx + 2\
({message_idx + 2}) > message_end_idx({message_end_idx}).")
                    break
                if (message_idx + trace_data_length) > message_end_idx:
                    print(
                        f"[WARNING] parse trace data message_idx + trace_data_length\
({message_idx + trace_data_length}) > message_end_idx({message_end_idx})."
                    )
                    break
                cur_data = data_list[message_idx : message_idx + trace_data_length]
                if added_data is not None:
                    cur_data += added_data
                trace_data_list.append(cur_data)
                message_idx += trace_data_length
            cur_idx = pack_end_idx
            message_cnt += 1
        frame_cnt = len(trace_data_list)
        print(f"message_cnt: {message_cnt}, frame_cnt: {frame_cnt}")
        return trace_data_list

    def is_task_blocked(self, state):
        """判断task是否blocked"""
        return (state == self.TASK_STATE_BLOCKED) or (state == self.TASK_STATE_SUSPEND)

    def get_task_state_text(self, state):
        """获取task的状态文本"""
        task_state_map = {
            self.TASK_STATE_RUNNING: "R",
            self.TASK_STATE_READY: "R",
            self.TASK_STATE_BLOCKED: "S",
            self.TASK_STATE_SUSPEND: "D",
            self.TASK_STATE_PENDING: "R",
            self.TASK_STATE_DELAYED: "T",
            self.TASK_STATE_DELETED: "X",
        }
        return task_state_map.setdefault(state, "S")

    def sort_trace_data_list(self, trace_data_list):
        """给二进制trace数据帧排序,按时间戳递增顺序"""
        trace_data_list.sort(key=lambda x: int.from_bytes(x[8:16], byteorder="little"))
        return trace_data_list

    def parse_payload_common(self, payload_list) -> dict:
        """解析下位机trace数据结构体的common部分"""
        priority: int = payload_list[2]
        core_id: int = payload_list[3]
        pid: int = int.from_bytes(payload_list[4:8], byteorder="little")
        tick: int = int.from_bytes(payload_list[8:16], byteorder="little")
        time_stamp: float = tick * self.tick_to_ns + self.time_offset

        return priority, core_id, pid, time_stamp

    def parse_irq_entry(self, event_len, payload_list):
        """解下位机结构体ftrace_irqentry_s"""
        if event_len != self.TRACE_IRQ_ENTRY_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_irq_entry(): \
length of event is not {self.TRACE_IRQ_ENTRY_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        irq_id: int = int.from_bytes(payload_list[16:20], byteorder="little")
        vector_pid: int = int.from_bytes(payload_list[20:24], byteorder="little")
        return [core_id, pid, time_stamp, irq_id, irq_id, vector_pid]

    def parse_irq_exit(self, event_len, payload_list):
        """解下位机结构体ftrace_irqexit_s"""
        if event_len != self.TRACE_IRQ_EXIT_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_irq_exit(): \
length of event is not {self.TRACE_IRQ_EXIT_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        irq_id: int = int.from_bytes(payload_list[16:20], byteorder="little")
        return [core_id, pid, time_stamp, irq_id, irq_id]

    def parse_sched_wakeup(self, event_len, payload_list):
        """解下位机结构体ftrace_wakeup_s"""
        if event_len != self.TRACE_SCHED_WAKEUP_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_sched_wakeup(): \
length of event is not {self.TRACE_SCHED_WAKEUP_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        wakeup_pid: int = int.from_bytes(payload_list[16:20], byteorder="little")
        return [priority, core_id, pid, time_stamp, wakeup_pid]

    def parse_sched_switch(self, event_len, payload_list):
        """解下位机结构体ftrace_schedswitch_s"""
        if event_len != self.TRACE_SCHED_SWITCH_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_sched_switch(): \
length of event is not {self.TRACE_SCHED_SWITCH_LEN}."
            )
            return None
        next_prioirty, core_id, next_pid, time_stamp = self.parse_payload_common(payload_list)
        prev_pid: int = int.from_bytes(payload_list[16:20], byteorder="little")
        prev_state: int = payload_list[20]
        return [next_prioirty, core_id, next_pid, time_stamp, prev_pid, prev_state]

    def parse_reactivate(self, event_len, payload_list):
        """解下位机结构体ftrace_reactivate_s"""
        if event_len != self.TRACE_REACTIVATE_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_reactivate(): \
length of 'reactivate' event is not {self.TRACE_REACTIVATE_LEN}."
            )
            return None
        return self.parse_payload_common(payload_list)

    def parse_event_start(self, event_len, payload_list):
        """解下位机结构体ftrace_event_s"""
        if event_len != self.TRACE_EVENT_START_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_event_start(): \
length of 'event_start' event is not {self.TRACE_EVENT_START_LEN}."
            )
            return None
        prioirty, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        event_mask: int = int.from_bytes(payload_list[16:24], byteorder="little")
        return [core_id, pid, time_stamp, event_mask]

    def parse_runnable_s(self, event_len, payload_list):
        """解下位机结构体ftrace_runnable_s"""
        if event_len != self.TRACE_RUNNABLE_S_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_runnable_s(): \
length of 'ftrace_runnable_s' event is not {self.TRACE_RUNNABLE_S_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        measurement_id: int = (payload_list[17] << 8) | payload_list[16]
        return [priority, core_id, pid, time_stamp, measurement_id]

    def parse_function_s(self, event_len, payload_list, is_start: bool):
        """解下位机结构体ftrace_function_s"""
        if event_len != self.TRACE_FUNCTION_S_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_function_s(): \
length of 'ftrace_function_s' event is not {self.TRACE_FUNCTION_S_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        measurement_id: int = (payload_list[17] << 8) | payload_list[16]
        # parse the str events
        if is_start is True:
            len = 18
            while payload_list[len] != 0:
                len = len + 1
            events = payload_list[18:len].decode()
        else:
            events = None
        return [core_id, pid, time_stamp, measurement_id, events]

    def parse_spinlock(self, event_len, payload_list):
        """解析下位机结构体ftrace_spinlock_s"""
        if event_len != self.TRACE_SPINLOCK_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_spinlock(): \
length of 'spinlock' event is not {self.TRACE_SPINLOCK_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        spinlock_id: int = (payload_list[17] << 8) | payload_list[16]
        return [core_id, time_stamp, spinlock_id]

    def parse_resource(self, event_len, payload_list):
        """解析下位机结构体ftrace_resource_s"""
        if event_len != self.TRACE_RESOURCE_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_resource(): \
length of 'resource' event is not {self.TRACE_RESOURCE_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        resource_id: int = int.from_bytes(payload_list[16:20], byteorder="little")
        return [core_id, time_stamp, resource_id]

    def parse_interrupt_s(self, event_len, payload_list):
        """解析下位机结构体ftrace_interrupt_s"""
        if event_len != self.TRACE_INTERRUPT_S_LEN:
            print(f"[ERROR] parse_interrupt_s(): length of 'interrupt_s' event is not {self.TRACE_INTERRUPT_S_LEN}.")
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        return [core_id, time_stamp]


class TraceVCOS200Parser:
    VERSION_STR = VersionStr.VCOS_200.value
    VERSION_NO = VersionNumber.VCOS_200.value
    # 报文长度
    TRACE_IRQ_LEN = 21
    TRACE_SCHED_WAKEUP_LEN = 20
    TRACE_SUSPEND_LEN = 17
    TRACE_RESUME_LEN = 16
    TRACE_REACTIVATE_LEN = 16
    # 报文类型字段
    FTRACE_TASK_SUSPEND = 2
    FTRACE_TASK_RESUME = 3
    FTRACE_WAKE_UP = 4
    FTRACE_TASK_REACTIVATE = 5
    FTRACE_IRQ_ENTRY = 22
    FTRACE_IRQ_EXIT = 23
    FTRACE_EVENT_START = -1
    # 状态值
    TASK_STATE_SUSPEND = 4  # INACTIVE

    def __init__(self):
        self.tick_to_ns: float = 0.0  # 1tick等于多少ns
        self.time_offset: int = 0  # 时间戳偏移(ns)

    def __getattr__(self, name):
        # 当访问不存在的属性时，返回 -1
        return -1

    def get_mapped_id(self, ori_id, core_id):
        """针对不同core的对象id进行映射"""
        return (ori_id << 4) + core_id

    def get_unmapped_id(self, mapped_id):
        return mapped_id >> 4

    def get_trace_data_list(self, data_list, added_data: bytearray = None):
        """将二进制数据分成trace数据帧"""
        trace_pack_start = bytearray([0xA5, 0xA5])
        trace_data_list: list = []
        total_len: int = len(data_list)
        cur_idx: int = 0
        last_seq_idx = -1
        max_seq_idx = (1 << 16) - 1
        pack_header_len = 12
        message_cnt = 0
        while cur_idx < total_len:
            # 每个报文包至少长12字节
            if (cur_idx + pack_header_len) > total_len:
                return trace_data_list
            # 包的开始字符
            if data_list[cur_idx : cur_idx + 2] != trace_pack_start:
                cur_idx += 1
                continue
            # 包的长度
            pack_len: int = (data_list[cur_idx + 5] << 8) | data_list[cur_idx + 4]
            pack_end_idx: int = cur_idx + pack_len
            # 当前包无效，找下一个包
            if (pack_end_idx > total_len) or (pack_len < pack_header_len):
                cur_idx += 1
                continue
            # 当前包有效,获取序列号
            cur_seq_idx = (data_list[cur_idx + 7] << 8) | data_list[cur_idx + 6]
            if last_seq_idx != -1:
                if last_seq_idx == max_seq_idx:
                    if cur_seq_idx != 0:
                        trace_data_list = []
                elif (cur_seq_idx - last_seq_idx) != 1:
                    trace_data_list = []
            last_seq_idx = cur_seq_idx
            # 获取每个数据帧
            message_idx: int = cur_idx + pack_header_len
            message_end_idx: int = pack_end_idx
            while message_idx < message_end_idx:
                trace_data_length = data_list[message_idx]
                if (message_idx + 2) > message_end_idx:
                    print(f"[WARNING] parse trace data message_idx + 2\
({message_idx + 2}) > message_end_idx({message_end_idx}).")
                    break
                if (message_idx + trace_data_length) > message_end_idx:
                    print(f"[WARNING] parse trace data message_idx + trace_data_length\
({message_idx + trace_data_length}) > message_end_idx({message_end_idx}).")
                    break
                if trace_data_length == 0:
                    print(f"[WARNING] A trace data's length value is 0.")
                    break
                cur_data = data_list[message_idx : message_idx + trace_data_length]
                if added_data is not None:
                    cur_data += added_data
                trace_data_list.append(cur_data)
                message_idx += trace_data_length
            cur_idx = pack_end_idx
            message_cnt += 1
        frame_cnt = len(trace_data_list)
        print(f"message_cnt: {message_cnt}, frame_cnt: {frame_cnt}")
        return trace_data_list

    def is_task_blocked(self, state):
        """判断task是否blocked"""
        if state == 0:
            return True
        elif state <= 3:
            return False
        return True

    def get_task_state_text(self, state):
        """获取task的状态文本"""
        if state == 0:
            return "X"
        elif state <= 3:
            return "R"
        return "S"

    def sort_trace_data_list(self, trace_data_list):
        """给二进制trace数据帧排序,按时间戳递增顺序"""
        trace_data_list.sort(key=lambda x: int.from_bytes(x[8:16], byteorder="little"))
        return trace_data_list

    def parse_payload_common(self, payload_list) -> dict:
        """解析下位机trace数据结构体的common部分"""
        priority: int = payload_list[2]
        core_id: int = payload_list[3]
        pid: int = int.from_bytes(payload_list[4:8], byteorder="little")
        pid = self.get_mapped_id(pid, core_id)
        tick: int = int.from_bytes(payload_list[8:16], byteorder="little")
        time_stamp: float = tick * self.tick_to_ns + self.time_offset
        return priority, core_id, pid, time_stamp

    def parse_irq(self, event_len, payload_list):
        """解note_irqhandler_s"""
        if event_len != self.TRACE_IRQ_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_irq(): length of ISR event is not {self.TRACE_IRQ_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        vector_pid: int = int.from_bytes(payload_list[16:20], byteorder="little")
        irq_id: int = payload_list[20]
        mapped_irq_id = self.get_mapped_id(irq_id, core_id)
        return [core_id, pid, time_stamp, mapped_irq_id, irq_id, vector_pid]

    def parse_irq_entry(self, event_len, payload_list):
        """irq进入"""
        return self.parse_irq(event_len, payload_list)

    def parse_irq_exit(self, event_len, payload_list):
        """irq结束"""
        result = self.parse_irq(event_len, payload_list)
        if result is None:
            return None
        core_id, pid, time_stamp, mapped_irq_id, irq_id, vector_pid = result
        return [core_id, pid, time_stamp, mapped_irq_id, irq_id]

    def parse_sched_wakeup(self, event_len, payload_list):
        """解note_wakeup_s"""
        if event_len != self.TRACE_SCHED_WAKEUP_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_sched_wakeup(): \
length of event is not {self.TRACE_SCHED_WAKEUP_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        wakeup_pid: int = int.from_bytes(payload_list[16:20], byteorder="little")
        wakeup_pid = self.get_mapped_id(wakeup_pid, core_id)
        return [priority, core_id, pid, time_stamp, wakeup_pid]

    def parse_suspend(self, event_len, payload_list):
        """解note_suspend_s"""
        if event_len != self.TRACE_SUSPEND_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_suspend(): \
length of 'note_suspend_s' is not {self.TRACE_SUSPEND_LEN}."
            )
            return None
        priority, core_id, pid, time_stamp = self.parse_payload_common(payload_list)
        task_state = payload_list[16]
        return [core_id, pid, time_stamp, task_state]

    def parse_resume(self, event_len, payload_list):
        """解note_resume_s"""
        if event_len != self.TRACE_RESUME_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_resume(): \
length of 'note_resume_s' is not {self.TRACE_RESUME_LEN}."
            )
            return None
        return self.parse_payload_common(payload_list)

    def parse_reactivate(self, event_len, payload_list):
        """解note_reactivate_s"""
        if event_len != self.TRACE_REACTIVATE_LEN:
            print(
                f"[ERROR] In VERSION({self.VERSION_STR}), parse_reactivate(): \
length of 'reactivate' event is not {self.TRACE_REACTIVATE_LEN}."
            )
            return None
        return self.parse_payload_common(payload_list)


TRACE_DATA_PARSER: Dict[int, Union[TraceVCOS100Parser, TraceVCOS200Parser]] = {
    VersionNumber.VCOS_100.value: TraceVCOS100Parser(),
    VersionNumber.VCOS_200.value: TraceVCOS200Parser(),
}
